##########################################################################################
# Title: Section 2 Parallel Driver → Distance Matrices + Figure 1 (Johannesburg)
# Section: 2 (Drivers)
# Purpose: Run Section 2 distance-matrix generation in parallel (by city) and produce Figure 1
# IO:
#   IN:  Processed_Data/City_Boundaries/<City>.gpkg (Mollweide), <City>_n.gpkg (native)
#        Data/Remote_Sensed/GHS_*_100m_YYYY.tif, Data/Remote_Sensed/NTL_2019.tif
#        Code/Function_Library/distance_matrix_bd.R
#   OUT: Processed_Data/Intermediate/Section_2/distance_matrix/<City>_<comb>.rds
#        Graphs/Figure_1.png (Johannesburg)
##########################################################################################

rm(list = ls()); options(stringsAsFactors = FALSE)

suppressPackageStartupMessages({
  library(terra)
  library(tidyterra) 
  library(patchwork)
  library(tmap)
  library(future)
  library(future.apply)
  library(ggplot2)
  library(data.table)
  library(sf)
})

# ---- Root and directories ----------------------------------------------------
if (!requireNamespace("here", quietly = TRUE)) install.packages("here", quiet = TRUE)
ROOT       <- normalizePath(here::here(), winslash = "/")
setwd(ROOT)

OUT_DM     <- file.path(ROOT, "Processed_Data", "Intermediate", "Section_2", "distance_matrix")
BOUNDARIES <- file.path(ROOT, "Processed_Data", "City_Boundaries")
GRAPHS_DIR <- file.path(ROOT, "Graphs")
dir.create(OUT_DM, recursive = TRUE, showWarnings = FALSE)
dir.create(GRAPHS_DIR, recursive = TRUE, showWarnings = FALSE)

# ---- Parallel + GDAL setup ---------------------------------------------------
WORKERS <- min(5, parallel::detectCores() - 1)
plan(multisession, workers = WORKERS)
terraOptions(progress = 0)
Sys.setenv(GDAL_NUM_THREADS = "ALL_CPUS")

# ---- Source core distance matrix function ------------------------------------
source(file.path(ROOT, "Code", "Function_Library", "distance_matrix_bd.R"))

# ---- Rasters -----------------------------------------------------------------
pick_raster_paths <- function(y) {
  if (y == 2020) return(list(pop = file.path(ROOT, "Data/Remote_Sensed/GHS_POP_100_2020.tif"),
                             bld = file.path(ROOT, "Data/Remote_Sensed/GHS_BUILT_100m_2020.tif")))
  if (y == 2010) return(list(pop = file.path(ROOT, "Data/Remote_Sensed/GHS_POP_100_2010.tif"),
                             bld = file.path(ROOT, "Data/Remote_Sensed/GHS_BUILT_100m_2010.tif")))
  if (y == 2000) return(list(pop = file.path(ROOT, "Data/Remote_Sensed/GHS_POP_100_2000.tif"),
                             bld = file.path(ROOT, "Data/Remote_Sensed/GHS_BUILT_100m_2000.tif")))
  stop("Unsupported year: ", y)
}
NTL_PATH <- file.path(ROOT, "Data/Remote_Sensed/NTL_2019.tif")

# ---- Cities and cutoffs ------------------------------------------------------
cities <- c(
  "Johannesburg","Capetown","Tshwane","Durban","NM_Bay",
  "Nairobi","Lagos","Abidjan","Cairo","London","Istanbul","Seoul",
  "SaoPaulo","KualaLumpur","Moscow","Mexico_City","Shenzhen",
  "Frankfurt","Zurich","Bogota","Tehran","Dubai","New_York","Mumbai","Buenos_Aires"
)
building_cutoffs <- c(80000, 90000, 100000, 110000)
cluster_sizes    <- c(4, 9, 14)

# ---- Boundary loader ---------------------------------------------------------
read_bounds <- function(city) {
  list(
    grid = vect(file.path(BOUNDARIES, paste0(city, ".gpkg"))),
    nat  = vect(file.path(BOUNDARIES, paste0(city, "_n.gpkg")))
  )
}

# ---- Skip logic --------------------------------------------------------------
should_skip <- function(city, year, cluster_size, bd_building_cutoff) {
  (city == "Mumbai" && year == 2000) ||
    ((city %in% c("Lagos","Mumbai","Buenos_Aires")) && cluster_size == 9  && bd_building_cutoff == 110000) ||
    ((city %in% c("Abidjan","Buenos_Aires"))        && cluster_size == 14 && bd_building_cutoff == 110000) ||
    ( city == "Mumbai"                               && cluster_size == 14 && bd_building_cutoff == 90000)  ||
    ((city %in% c("Lagos","Mumbai"))                 && cluster_size == 14 && bd_building_cutoff %in% c(100000,110000))
}

# ---- Core runner for 2020 (parallel) -----------------------------------------
run_city_2020 <- function(city) {
  on.exit({ terra::tmpFiles(current = TRUE, remove = TRUE); gc() }, add = TRUE)
  y <- 2020
  rs_paths <- pick_raster_paths(y)
  grid <- vect(file.path(BOUNDARIES, paste0(city, ".gpkg")))
  nat  <- vect(file.path(BOUNDARIES, paste0(city, "_n.gpkg")))
  
  pop_r <- rast(rs_paths$pop)
  bld_r <- rast(rs_paths$bld)
  ntl_r <- rast(NTL_PATH)
  
  for (bd_cut in building_cutoffs) {
    for (cs in cluster_sizes) {
      if (should_skip(city, y, cs, bd_cut)) next
      invisible(distance_matrix_bd(
        pop_raster  = pop_r,
        building_raster = bld_r,
        ntl_raster  = ntl_r,
        city_sv     = grid,
        city_sv_n   = nat,
        sparse_area_pop_cutoff = 10,
        sparse_area_building_cutoff = 50000,
        bd_pop_cutoff = 50,
        bd_building_cutoff = bd_cut,
        cluster_distance = 2000,
        cluster_size = cs,
        create_map = 0,
        year = y,
        city_name = city,
        write_outputs = TRUE,
        out_dir = OUT_DM,
        file_format = "rds"
      ))
    }
  }
  rm(pop_r, bld_r, ntl_r, grid, nat)
  invisible(TRUE)
}

# ---- Compact fixed param runner (for 2000/2010) ------------------------------
run_city_fixed <- function(city, y, bd_cut = 100000, cs = 9) {
  on.exit({ terra::tmpFiles(current = TRUE, remove = TRUE); gc() }, add = TRUE)
  rs_paths <- pick_raster_paths(y)
  grid <- vect(file.path(BOUNDARIES, paste0(city, ".gpkg")))
  nat  <- vect(file.path(BOUNDARIES, paste0(city, "_n.gpkg")))
  pop_r <- rast(rs_paths$pop)
  bld_r <- rast(rs_paths$bld)
  ntl_r <- rast(NTL_PATH)
  
  if (!should_skip(city, y, cs, bd_cut)) {
    invisible(distance_matrix_bd(
      pop_raster  = pop_r,
      building_raster = bld_r,
      ntl_raster  = ntl_r,
      city_sv     = grid,
      city_sv_n   = nat,
      sparse_area_pop_cutoff = 10,
      sparse_area_building_cutoff = 50000,
      bd_pop_cutoff = 50,
      bd_building_cutoff = bd_cut,
      cluster_distance = 2000,
      cluster_size = cs,
      create_map = 0,
      year = y,
      city_name = city,
      write_outputs = TRUE,
      out_dir = OUT_DM,
      file_format = "rds"
    ))
  }
  rm(pop_r, bld_r, ntl_r, grid, nat)
  invisible(TRUE)
}

# ---- Parallel runs for all years --------------------------------------------
invisible(future_lapply(cities, run_city_2020))
invisible(future_lapply(cities, function(cty) run_city_fixed(cty, y = 2010)))
invisible(future_lapply(cities, function(cty) run_city_fixed(cty, y = 2000)))

# ---- Build Figure 1 (Johannesburg) ------------------------------------------
raster_maps_gg <- function(raster_s_actual, raster_s, city, boundary, out_path) {
  bld   <- mask(crop(raster_s_actual[["building"]],  boundary), boundary)
  bdcat <- mask(crop(raster_s_actual[["bd_dummy"]],  boundary), boundary)
  distp <- mask(crop(raster_s[["min_dist_cbd"]],     boundary), boundary)
  popp  <- mask(crop(raster_s[["pop"]],              boundary), boundary)
  
  # Fill NA inside boundary with zeros for rasters we plot
  bmask <- rasterize(boundary, bld, field = 1)
  bld0  <- ifel(is.na(bld)   & !is.na(bmask), 0, bld)
  popp0 <- ifel(is.na(popp)  & !is.na(bmask), 0, popp)
  bd0   <- ifel(is.na(bdcat) & !is.na(bmask), 0, bdcat)
  
  # Points for BD pixels == 1
  bd_pts_tv <- as.points(ifel(bd0 == 1, 1, NA), values = FALSE, na.rm = TRUE)
  bd_pts    <- sf::st_as_sf(bd_pts_tv)  # already POINT
  
  # Grey "no data" mask for distance panel
  na_mask_c   <- ifel(is.na(distp) & !is.na(bmask), 1, NA)
  na_polys_c  <- as.polygons(na_mask_c, dissolve = TRUE, na.rm = TRUE)
  na_polys_sf <- sf::st_as_sf(na_polys_c)
  
  # Bounds & background
  bnd_sf  <- sf::st_as_sf(boundary)
  bb_sfc  <- sf::st_as_sfc(sf::st_bbox(bnd_sf))
  xlim    <- sf::st_bbox(bb_sfc)[c("xmin","xmax")]
  ylim    <- sf::st_bbox(bb_sfc)[c("ymin","ymax")]
  bb_layer <- geom_sf(data = bb_sfc, fill = "#D9D9D9", color = NA)
  
  # Boundary halo (sf, not tidyterra)
  boundary_halo <- list(
    geom_sf(data = bnd_sf, fill = NA, color = "white", linewidth = 0.55),
    geom_sf(data = bnd_sf, fill = NA, color = "black", linewidth = 0.30)
  )
  
  # Theme + legend tweaks: smaller, more right, breathing room
  theme_base <- theme_minimal(base_size = 11) +
    theme(
      panel.background = element_rect(fill = "white", colour = NA),
      plot.background  = element_rect(fill = "white", colour = NA),
      panel.grid = element_blank(),
      legend.text  = element_text(size = 7.5),
      legend.title = element_text(size = 8)
    )
  
  legend_in_br <- theme(
    legend.position = "inside",
    legend.position.inside = c(0.992, 0.03),  # farther right
    legend.justification = c(1, 0),
    legend.background = element_blank(),
    legend.key.height = grid::unit(8,  "pt"),
    legend.key.width  = grid::unit(9,  "pt"),
    legend.box.margin = margin(0, 2, 0, 0)
  )
  
  theme_strip_axes <- theme(
    axis.title = element_blank(),
    axis.text  = element_blank(),
    axis.ticks = element_blank()
  )
  
  panel_margin <- theme(plot.margin = margin(t = 3, r = 12, b = 3, l = 3))
  
  guide_fill_small  <- guides(
    fill = guide_colorbar(
      barheight = grid::unit(46, "pt"),
      barwidth  = grid::unit(10, "pt"),
      ticks = FALSE,
      label.position = "right",
      title.position = "top"
    )
  )
  
  guide_shape_small <- guides(
    shape = guide_legend(
      title = NULL,
      override.aes = list(size = 2.0)
    )
  )
  
  # ---- Panels ----
  pA <- ggplot() + bb_layer + geom_spatraster(data = bld0, na.rm = TRUE) + boundary_halo +
    scale_fill_distiller(palette = "Reds", direction = 1, na.value = NA,
                         name = expression(10^3~m^3),
                         breaks = c(0, 1e5, 2e5), labels = c("0", "100", "200")) +
    labs(title = "A: Building Volumes") +
    coord_sf(xlim = xlim, ylim = ylim, expand = FALSE, clip = "off") +
    theme_base + legend_in_br + theme_strip_axes + panel_margin + guide_fill_small
  
  pB <- ggplot() + bb_layer + boundary_halo +
    geom_sf(data = bd_pts, aes(shape = "Business Districts"),
            color = "#C00000", alpha = 0.8, size = 0.01) +
    scale_shape_manual(values = c("Business Districts" = 20)) +
    labs(title = "B: Business Districts") +
    coord_sf(xlim = xlim, ylim = ylim, expand = FALSE, clip = "off") +
    theme_base + legend_in_br + theme_strip_axes + panel_margin + guide_shape_small
  
  reds_light <- RColorBrewer::brewer.pal(9, "Reds")[1:9]
  pC <- ggplot() + bb_layer +
    geom_sf(data = na_polys_sf, fill = "#C7C7C7", color = NA, alpha = 0.6) +
    geom_spatraster(data = distp, na.rm = TRUE) + boundary_halo +
    scale_fill_gradientn(colours = reds_light, limits = c(0, 100),
                         breaks = c(25, 50, 75, 100), na.value = NA, name = "Percentiles") +
    labs(title = "C: Distance to Business District") +
    coord_sf(xlim = xlim, ylim = ylim, expand = FALSE, clip = "off") +
    theme_base + legend_in_br + theme_strip_axes + panel_margin + guide_fill_small
  
  pD <- ggplot() + bb_layer + geom_spatraster(data = popp0, na.rm = TRUE) + boundary_halo +
    scale_fill_distiller(palette = "Reds", direction = 1, limits = c(0, 100),
                         breaks = c(0, 25, 50, 75, 100), na.value = NA, name = "Percentiles") +
    labs(title = "D: Population Density") +
    coord_sf(xlim = xlim, ylim = ylim, expand = FALSE, clip = "off") +
    theme_base + legend_in_br + theme_strip_axes + panel_margin + guide_fill_small
  
  final_plot <- (pA | pB) / (pC | pD)
  ggsave(out_path, plot = final_plot, width = 7, height = 8, dpi = 300, bg = "white")
  message("Saved: ", out_path)
}

# ---- Generate Figure 1 (Johannesburg) ---------------------------------------
jb <- read_bounds("Johannesburg")
dt <- as.data.table(readRDS(file.path(OUT_DM, "Johannesburg_100K_09_2020.rds")))
N <- nrow(dt); den <- max(1L, N - 1L)
dt[, pop_pct  := 100 * (frank(pop,          ties.method = "average") - 1) / den]
dt[, dist_pct := 100 * (frank(min_dist_cbd, ties.method = "average") - 1) / den]

make_stack <- function(df, cols, crs_str) {
  rs <- lapply(cols, function(cl) terra::rast(df[, .(x, y, z = get(cl))], type = "xyz", crs = crs_str))
  out <- do.call(c, rs); names(out) <- cols; out
}
raster_s_actual <- make_stack(dt, c("building","bd_dummy"), terra::crs(jb$nat))
raster_s        <- make_stack(dt, c("dist_pct","pop_pct"),  terra::crs(jb$nat))
names(raster_s) <- c("min_dist_cbd","pop")

raster_maps_gg(raster_s_actual, raster_s, "Johannesburg", jb$nat,
               file.path(GRAPHS_DIR, "Figure_1.png"))

message("✅ Section 2 parallel driver complete.")
